package com.xiashitech.agent.instrumentation.db.resultset.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.instrumentation.db.common.XSResultSizeContext;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSResultSetCallDepth;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.instrumentation.jdbc.internal.JdbcData;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;

import static com.xiashitech.agent.instrumentation.db.resultset.XSResultSetSingletons.instrumenter;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSResultSetInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("java.sql.ResultSet");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return implementsInterface(named("java.sql.ResultSet"));
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                named("next")
                        .and(takesNoArguments())
                        .and(returns(boolean.class))
                        .and(isPublic()),
                XSResultSetInstrumentation.class.getName() + "$XSResultLineAdvice");
        transformer.applyAdviceToMethod(
                nameStartsWith("get")
                        .and(takesArguments(1).or(takesArguments(2)))
                        .and(isPublic()),
                XSResultSetInstrumentation.class.getName() + "$XSResultColumnAdvice");
        transformer.applyAdviceToMethod(
                named("close")
                        .and(takesNoArguments())
                        .and(isPublic()),
                XSResultSetInstrumentation.class.getName() + "$XSResultCloseAdvice");
    }

    @SuppressWarnings("unused")
    public static class XSResultLineAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
            @Advice.This ResultSet resultSet,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(resultSet != null && resultSet.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                callDepth = CallDepth.forClass(XSResultSetCallDepth.class);
                if (callDepth.getAndIncrement() > 0) {
                    return;
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("XSResultLineAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
            @Advice.This ResultSet resultSet,
            @Advice.Return boolean columnResult,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(resultSet != null && resultSet.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                if (callDepth.decrementAndGet() > 0) {
                    return;
                }
                if(columnResult == true) {
                    XSResultSizeContext.setIsRecord(resultSet,true);
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("XSResultLineAdvice e");
        }
    }

    @SuppressWarnings("unused")
    public static class XSResultColumnAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
            @Advice.This ResultSet resultSet,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(resultSet != null && resultSet.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                callDepth = CallDepth.forClass(XSResultSetCallDepth.class);
                if (callDepth.getAndIncrement() > 0) {
                    return;
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("XSResultColumnAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
            @Advice.This ResultSet resultSet,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(resultSet != null && resultSet.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                if (callDepth.decrementAndGet() > 0) {
                    return;
                }
                Boolean isRecord = XSResultSizeContext.getIsRecord(resultSet);
                if(isRecord) {
                    XSResultSizeContext.record(resultSet);
                }
                XSResultSizeContext.setIsRecord(resultSet,false);
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("XSResultColumnAdvice e");
        }
    }

    @SuppressWarnings("unused")
    public static class XSResultCloseAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
            @Advice.This ResultSet resultSet,
            @Advice.Local("otelSql") String sql,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(resultSet != null && resultSet.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                callDepth = CallDepth.forClass(XSResultSetCallDepth.class);
                if (callDepth.getAndIncrement() > 0) {
                    return;
                }
                try {
                    if (!resultSet.isClosed() && resultSet.getStatement() != null && resultSet.getStatement() instanceof PreparedStatement) {
                        sql = JdbcData.preparedStatement.get((PreparedStatement)resultSet.getStatement());
                    } else if(!resultSet.isClosed() && resultSet.getStatement() != null && resultSet.getStatement() instanceof Statement) {
                        XSResultSizeContext.StatementSQLContext statementSQLContext = XSResultSizeContext.getStatementSQLContext(resultSet.getStatement());
                        if(statementSQLContext != null) {
                            sql = statementSQLContext.getSql();
                        }
                    }
                } catch (Throwable ee) {
                    StartConfig.throwableStackPrinter(ee);
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("XSResultCloseAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
            @Advice.This ResultSet resultSet,
            @Advice.Local("otelSql") String sql,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(resultSet != null && resultSet.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                if (callDepth.decrementAndGet() > 0) {
                    return;
                }
                try {
                    XSResultSizeContext.ResultSetContext resultSetContext = XSResultSizeContext.getResultSetContext(resultSet);
                    if(resultSetContext != null) {
                        Context context = resultSetContext.getContext();
                        Throwable throwable = resultSetContext.getThrowable();
                        XSResultSizeContext.Request request = XSResultSizeContext.Request.builder().resultSet(resultSet).size(XSResultSizeContext.getSize(resultSet)).sql(sql).build();
                        instrumenter().end(context, request, null, throwable);
                    }
                } catch (Throwable ee) {
                    StartConfig.throwableStackPrinter(ee);
                } finally {
                    XSResultSizeContext.refreshCurrentThreadLocal();
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("XSResultCloseAdvice e");
        }
    }
}
